# coding=utf-8
import os
import shutil
import threading
from queue import Queue

import pandas as pd
from loguru import logger

from cal_ops.point import cal_point1
from mylib import download_all
from mylib.mydb import MyDB

st_queue = Queue()


def clear_dir(d='result'):
    if os.path.exists(d):
        shutil.rmtree(d)
    os.makedirs(d)


def get_all_stock_csv_path():
    for root, dirs, files in os.walk('stocks'):
        return [os.path.join(root, item) for item in files]


# def change_all_stockcsv2sqlite():
#     all_stocks_path_list = get_all_stock_csv_path()
#     for sp in all_stocks_path_list:
#         table_name = sp.replace('.csv', '').replace('stocks/', 'stock_').replace('.', '_')
#         mydb.change_csv2sqlite(csv_name=sp, table_name=table_name, from_col=1)


def get_all_stocks_indus():
    with open('industry_count_rank.csv', 'r') as fr:
        all_indus = [item.strip().split(',') for item in fr.readlines()[1:]]
        return all_indus


def get_avg_top10(stocks, N=50):
    stocks_top_dict = {}
    for stock in stocks:
        csv_path = f'stocks/{stock}.csv'
        if not os.path.exists(csv_path):
            continue
        df2 = pd.read_csv(csv_path)
        stocks_top_dict[stock] = round(abs(df2[0:N]['pct_chg'].mean()), 4)
    sorted_d = sorted(stocks_top_dict.items(), key=lambda x: x[1], reverse=True)
    # 获取近N日振幅排名前10的股票
    return sorted_d[:10]


def get_end_date(log_dir, indus, indus_num):
    f_path = f'{log_dir}/{indus_num}_all_{indus}.csv'
    if not os.path.exists(f_path):
        return None
    with open(f_path, 'r') as fr:
        line = fr.readline()
        return eval(line.split(',')[0])


def t_download():
    all_indus = get_all_stocks_indus()
    for indus, indus_num in all_indus:
        # 快速调试用
        # if eval(indus_num) >= 5:
        #     break
        mydb_obj = MyDB()
        stocks = mydb_obj.select_stocks_from_db_by_industry(indus)
        mydb_obj.close()
        logger.warning(f't_download {indus} {indus_num}')
        stocks = download_all.run(stocks)
        # if len(stocks) > 10:
        #     stocks = get_avg_top10(stocks)
        #     stocks = [s for s, p in stocks]
        st_queue.put((indus, indus_num, stocks))
    st_queue.put((1, 1, 1))


def t_cal():
    log_dir = os.path.basename(__file__).split('.')[0]
    # 需要个股涨跌图打开，一般只计算一天才打开
    gen_excel = False
    bkd = 1000
    start_date = 20301230
    while 1:
        if not st_queue:
            continue
        indus, indus_num, stocks = st_queue.get()
        if (1, 1, 1) == (indus, indus_num, stocks):
            logger.warning('所有计算完成')
            break
        logger.warning(f't_cal 正在计算【{indus}, {indus_num}】')
        end_date = get_end_date(log_dir, indus, indus_num)
        cal_point1.run(gen_excel, log_dir, start_date, end_date, stocks, indus_num, indus=indus, bkd=bkd, N=6)
        logger.warning(f't_cal 计算完成【{indus}, {indus_num}】')


if __name__ == '__main__':
    logger.add("log/main_all_indus_{time}.log", level='WARNING')
    t1 = threading.Thread(target=t_download, args=())
    t2 = threading.Thread(target=t_cal, args=())
    t1.start()
    t2.start()
    t1.join()
    t2.join()
